//
// Created by myrica on 4/25/25.
//

#ifndef MOVEALLOCALLVMIR_H
#define MOVEALLOCALLVMIR_H

#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Passes/PassBuilder.h"



using namespace llvm;

namespace {
    struct MoveAllocaToEntryPass : public PassInfoMixin<MoveAllocaToEntryPass> {
        PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
            BasicBlock &Entry = F.getEntryBlock();
            Instruction *InsertPt = &*Entry.getFirstInsertionPt();

            std::vector<AllocaInst *> AllocasToMove;

            for (BasicBlock &BB : F) {
                for (Instruction &I : BB) {
                    if (auto *AI = dyn_cast<AllocaInst>(&I)) {
                        if (AI->getParent() != &Entry) {
                            AllocasToMove.push_back(AI);
                        }
                    }
                }
            }

            for (AllocaInst *AI : AllocasToMove) {
                AI->moveBefore(InsertPt);
            }

            return PreservedAnalyses::all();
        }
    };
}


llvm::PassPluginLibraryInfo getMoveAllocaToEntryPassPluginInfo();

extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
    return getMoveAllocaToEntryPassPluginInfo();
}

#endif //MOVEALLOCALLVMIR_H
